# 首先，确保安装了必要的库
# 如果没有安装，可以使用以下命令进行安装：
# pip install akshare pandas matplotlib numpy

import akshare as ak
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import datetime

def get_stock_data(stock_code, start_date, end_date):
    """
    获取股票数据
    :param stock_code: 股票代码，例如'000001.SZ'表示平安银行
    :param start_date: 开始日期，格式'YYYY-MM-DD'
    :param end_date: 结束日期，格式'YYYY-MM-DD'
    :return: pandas DataFrame
    """
    # 使用akshare的stock_zh_a_daily函数获取日线数据
    df = ak.stock_zh_a_daily(symbol=stock_code, start_date=start_date, end_date=end_date, adjust="qfq")
    # 将日期列转换为datetime类型，并设置为索引
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)
    # 按日期升序排序
    df.sort_index(inplace=True)
    return df

def initialize_grid(df, grid_spread, grid_levels, initial_investment):
    """
    初始化网格
    :param df: 股票数据
    :param grid_spread: 网格间距（百分比）
    :param grid_levels: 网格层级数
    :param initial_investment: 初始投资金额
    :return: 网格数据框
    """
    # 初始价格
    initial_price = df['close'].iloc[0]
    # 生成网格价格
    grid_prices = []
    for i in range(-grid_levels, grid_levels + 1):
        price = initial_price * (1 + i * grid_spread)
        grid_prices.append(price)
    grid_prices = sorted(grid_prices)
    # 创建网格数据框
    grid = pd.DataFrame({
        'price': grid_prices,
        'action': ['sell'] * grid_levels + ['hold'] + ['buy'] * grid_levels,
        'position': 0
    })
    # 计算初始持仓
    initial_shares = initial_investment / initial_price
    grid.loc[grid['action'] == 'sell', 'position'] = initial_shares / (2 * grid_levels)
    return grid

def execute_trades(df, grid):
    """
    执行交易逻辑
    :param df: 股票数据
    :param grid: 网格数据框
    :return: 回测结果
    """
    cash = 100000  # 初始现金10万元
    shares = 0
    portfolio = []
    for idx, row in df.iterrows():
        price = row['close']
        # 检查是否达到买入或卖出点位
        for i in range(len(grid)):
            if grid['action'].iloc[i] == 'buy' and price <= grid['price'].iloc[i] and grid['position'].iloc[i] > 0:
                # 执行买入
                buy_shares = grid['position'].iloc[i]
                cash -= buy_shares * price
                shares += buy_shares
                grid.at[i, 'position'] = 0  # 重置网格
                print(f"{idx} 买入 {buy_shares} 股，价格 {price}，剩余现金 {cash}")
            elif grid['action'].iloc[i] == 'sell' and price >= grid['price'].iloc[i] and grid['position'].iloc[i] > 0:
                # 执行卖出
                sell_shares = grid['position'].iloc[i]
                cash += sell_shares * price
                shares -= sell_shares
                grid.at[i, 'position'] = 0  # 重置网格
                print(f"{idx} 卖出 {sell_shares} 股，价格 {price}，剩余现金 {cash}")
        # 更新持仓
        portfolio.append({'date': idx, 'price': price, 'shares': shares, 'cash': cash, 'total': cash + shares * price})
    return pd.DataFrame(portfolio)

def plot_results(df, portfolio):
    """
    绘制回测结果
    :param df: 股票数据
    :param portfolio: 回测结果
    """
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10), sharex=True)

    # 绘制价格和网格线
    ax1.plot(df.index, df['close'], label='收盘价', color='blue')
    for i in range(len(grid)):
        ax1.axhline(y=grid['price'].iloc[i], color='red', linestyle='--', linewidth=0.5)
    ax1.set_title('网格交易策略回测')
    ax1.set_ylabel('价格')
    ax1.legend()

    # 绘制总资产
    ax2.plot(portfolio['date'], portfolio['total'], label='总资产', color='b')
    ax2.set_ylabel('总资产 (元)')
    ax2.legend()

    plt.xlabel('日期')
    plt.show()


def test():
    # 网格交易参数设置
    symbol = "sz000011"  # 平安银行的股票代码作为例子
    start_date = "20240101"
    end_date = "20250101"
    initial_balance = 10000  # 初始资金
    grid_size = 15.0  # 每个网格的大小（百分比）
    num_grids = 1000  # 网格的数量
    buy_volume_per_grid = 100  # 每个网格购买的股数

    # 获取历史行情数据
    stock_data = ak.stock_zh_a_daily(symbol=symbol, start_date=start_date, end_date=end_date, adjust="qfq")

    # 准备回测环境
    balance = initial_balance
    holdings = 0
    entry_price = None
    positions = []
    sum = 0 # 当前总收益
    total_value = []  # 用于存储每个时间点的投资组合价值

    # 计算网格价格水平
    print(stock_data.columns.to_list())
    min_price = stock_data['close'].min()
    max_price = stock_data['close'].max()
    price_levels = [min_price + i * (max_price - min_price) / num_grids for i in range(num_grids + 1)]

    # 回测逻辑
    for index, row in stock_data.iterrows():
        close_price = row['close']

        # 卖出逻辑：当持有仓位且当前价格高于进入价格加上一个网格大小时卖出
        if holdings > 0 and entry_price is not None:
            upper_level = entry_price * (1 + grid_size / 100)
            if close_price >= upper_level:
                balance += holdings * close_price
                holdings = 0
                entry_price = None
                positions.append(('Sell', close_price))
                print(f"Sell at {close_price}")

        # 买入逻辑：当没有持仓且当前价格低于网格价格水平时买入
        elif balance > close_price * buy_volume_per_grid:
            for level in price_levels:
                if close_price <= level:
                    balance -= close_price * buy_volume_per_grid
                    holdings += buy_volume_per_grid
                    entry_price = close_price
                    positions.append(('Buy', close_price))
                    print(f"Buy at {close_price}")
                    break

        total_value.append(holdings * close_price + balance)

    # 输出最终结果
    final_value = balance + holdings * stock_data.iloc[-1]['close']
    print(f"Initial Balance: {initial_balance}, Final Value: {final_value}")

    plt.figure(figsize=(14, 7))
    plt.subplot(2, 1, 1)
    plt.plot(stock_data['close'], label='Close Price', color='blue')
    plt.title('Stock Close Price Over Time')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.grid(True)

    # 绘制收益曲线（累计收益率）
    plt.subplot(2, 1, 2)
    plt.plot(total_value, label='Cumulative Return', color='green')
    plt.title('Cumulative Returns Over Time')
    plt.xlabel('Date')
    plt.ylabel('Cumulative Return')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.show()

def main():
    # 参数设置
    stock_code = 'sz000001'  # 平安银行
    start_date = '20200101'
    end_date = '20250101'
    grid_spread = 0.02  # 网格间距2%
    grid_levels = 5  # 网格层级5
    initial_investment = 100000  # 初始投资10万元

    # 获取数据
    df = get_stock_data(stock_code, start_date, end_date)
    print(f"获取到的数据日期范围: {df.index.min()} 到 {df.index.max()}")

    # 初始化网格
    grid = initialize_grid(df, grid_spread, grid_levels, initial_investment)
    print("网格点位:")
    print(grid)

    # 执行交易
    portfolio = execute_trades(df, grid)

    # 输出回测结果
    total_return = (portfolio['total'].iloc[-1] / initial_investment) - 1
    print(f"初始投资: {initial_investment}元")
    print(f"最终总资产: {portfolio['total'].iloc[-1]}元")
    print(f"总收益率: {total_return * 100:.2f}%")

    # 绘制图表
    plot_results(df, portfolio)

if __name__ == "__main__":
    test()